///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/davidson.h
///
/// \brief   header file for the Davidson solver
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef FORMIC_DAVIDSON_HEADER
#define FORMIC_DAVIDSON_HEADER

#include<cmath>
#include<algorithm>
#include<complex>

#include<formic/exception.h>
#include<formic/matrix_actor.h>
#include<formic/numeric/numeric.h>
#include<formic/mpi/interface.h>
#include<formic/lapack/interface.h>
#include<formic/output/output.h>
#include<formic/random/random.h>

namespace formic {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  Solves a generalized eigenvalue problem  ( H x = E S x )  for the
  ///         (eigenvalue,eigenvector) pair (E,x) with the lowest eigenvalue using
  ///         the iterative Davidson solver.
  ///         
  ///         The solver constructs the matrices H and S in a subspace spanned by the vectors
  ///         (v_0, v_1, ...) and solves for the lowest eigenpair in this subspace.
  ///
  ///         The subpace is expanded iteratively by adding the vector
  ///
  ///         v_(i+1) = M ( H - E_i S ) x_i
  ///
  ///         where (E_i,x_i) is the current approximation to the desired eigenpair and M
  ///         is a preconditioning matrix.
  ///
  ///         Typically, this preconditioning matrix is chosen to be
  ///
  ///         ( diag(H) - E_i diag(S) )^(-1)
  ///
  ///         where diag(A) means the matrix with diagonal elements from A and zero off-diagonals.
  ///
  ///         In general, however, the matrix M may be chosen to be any matrix that is a good
  ///         approximation to the inverse matrix pencil ( H - E_i S )^(-1).
  ///
  /// \param[in]      n           the dimension of the matrices and vectors
  /// \param[in]      m           the maximum number of iterations
  /// \param[in]      thresh      threshold for convergence
  /// \param[in,out]  E           on entry, a guess for the eigenvalue.
  ///                             on exit, the subspace-approximated eigenvalue.
  /// \param[in,out]  x           on entry, a guess for the eigenvector.
  ///                             on exit, the subspace-approximated eigenvector.
  /// \param[in,out]  H           object that multiplies by the matrix H
  /// \param[in,out]  S           object that multiplies by the Hermitian, positive definite matrix S
  /// \param[in,out]  M           object that multiplies by the matrix M
  /// \param[in]      symmetrize  whether to force symmetrization of the subspace H matrix
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////

  template<class T> inline void davidson(const int n,
                                         const int m,
                                         const double thresh,
                                         double & E,
                                         T * const x,
                                         formic::MatrixActor<T> * const H,
                                         formic::MatrixActor<T> * const S,
                                         formic::MatrixActor<T> * const M,
                                         const bool symmetrize) {

    // make sure the dimension is sane
    if (n <= 0)
      throw formic::Exception("invalid matrix dimension of \"%i\" in formic::davidson") % n;

    // make sure the number of iterations is sane
    if (m <= 0)
      throw formic::Exception("\"%i\" is not a valid number of iterations in formic::davidson") % m;

    // make sure the threshold is sane
    if (thresh <= 0.0)
      throw formic::Exception("the threshold must be positive in formic::davidson");

    // initialize the residual vector
    std::vector<T> w(n);

    // initialize vector to hold the best solution
    std::vector<T> best_x(n);
    formic::xcopy(n, x, 1, &best_x.at(0), 1);

    // create a vector to hold the subspace basis vectors (v_0, v_1, ...)
    std::vector<std::vector<T> > v(1+m);
    v.at(0).resize(n);

    // set the initial guess as the first basis vector
    formic::xcopy(n, x, 1, &v.at(0).at(0), 1);

    // initialize the subspace matrices
    std::vector<T> subH;
    std::vector<T> subS;

    // create vectors of vectors to hold matrix-vector products
    std::vector<std::vector<T> > Hv(m);
    std::vector<std::vector<T> > Sv(m);

    // create arrays and other variables used in diagonalization routines
    std::vector<T> matA(m*m, formic::zero(T()));
    std::vector<T> matB(m*m, formic::zero(T()));
    std::vector<double> wr(m, 0.0);
    std::vector<double> wi(m, 0.0);
    std::vector<T> wc(m, formic::zero(T()));
    int lwork = (10 + m) * m;
    std::vector<T> work(lwork, formic::zero(T()));
    std::vector<double> rwork(lwork, 0.0);
    int info = 0;

    // create arrays to hold the S matrix eigenvalues and eigenvectors
    std::vector<double> s_evals(m, 0.0);
    std::vector<T> s_evecs(m*m, formic::zero(T()));

    // create an array to hold the coefficients for the eigenvector in the subspace
    std::vector<T> sub_evec(m, formic::zero(T()));

    // initialize the residual norm
    double rnorm = 1.0e100;

    // initialize a variable to hold the previous iteration's eigenvalue
    double last_eval = 1.0e100;

    // iteratively expand the subspace matrices and basis vectors
    int iter = 0;
    bool converged = false;
    for ( ; iter < m && !converged; iter++) {

      // get the new subspace dimension
      const int p = iter + 1;

      //// pre-normalize the new basis vector
      //{
      //  S->operate(n, &v.at(iter).at(0), x, formic::mpi::rank(), 0);
      //  const T vnorm = std::sqrt(formic::xdotc(n, &v.at(iter).at(0), 1, x, 1));
      //  formic::xscal(n, formic::unity(T()) / vnorm, &v.at(iter).at(0), 1);
      //}

      //// add random noise to new vector
      //for (int i = 0; i < n; i++)
      //  v.at(iter).at(i) += ( 1.0e-3 / n ) * ( formic::random_number<T>() - 0.5 * formic::unity(T()) );

      //// orthogonalize the new vector against the previous vectors
      //for (int i = 0; i < iter; i++)
      //  formic::xaxpy(n, -formic::xdotc(n, &Sv.at(i).at(0), 1, &v.at(iter).at(0), 1), &v.at(i).at(0), 1, &v.at(iter).at(0), 1);

      // compute the product of S with the new (un-normalized) basis vector and place it temporarily in x
      S->operate(n, &v.at(iter).at(0), x, formic::mpi::rank(), 0);

      // normalize the new basis vector
      double vnorm = std::abs(std::sqrt(formic::xdotc(n, &v.at(iter).at(0), 1, x, 1)));
      formic::mpi::bcast(vnorm);
      formic::xscal(n, formic::unity(T()) / vnorm, &v.at(iter).at(0), 1);

      //// print the basis vector
      //formic::of << boost::format("basis vector %2i  (norm = %20.12f):") % iter % std::abs(vnorm) << std::endl;
      //for (int i = 0; i < n; i++)
      //  formic::of << formic::format_number("%20.12f", v.at(iter).at(i)) << std::endl;
      //formic::of << std::endl;

      // place the product of S with the new (normalized) basis vector in the Sv vector
      Sv.at(iter).assign(n, formic::zero(T()));
      formic::xaxpy(n, formic::unity(T()) / vnorm, x, 1, &Sv.at(iter).at(0), 1);

      // compute the product of H with the new (normalized) basis vector and store it in Hv
      Hv.at(iter).resize(n);
      H->operate(n, &v.at(iter).at(0), &Hv.at(iter).at(0), formic::mpi::rank(), 0);

      // create new matrices for the expanded subspace
      {

        // allocate temporary matrices in which to build the new matrices
        std::vector<T> newH(p*p, formic::zero(T()));
        std::vector<T> newS(p*p, formic::zero(T()));

        // copy the old matrices into the new matrices' upper left block
        for (int i = 0; i < iter; i++)
        for (int j = 0; j < iter; j++) {
          newH.at( i + p*j ) = subH.at( i + iter*j );
          newS.at( i + p*j ) = subS.at( i + iter*j );
        }

        // compute the new row and column of the subspace matrices
        for (int i = 0; i < p; i++) {
          newH.at( i    + p*iter ) = formic::xdotc(n, &v.at(i   ).at(0), 1, &Hv.at(iter).at(0), 1);
          newH.at( iter + p*i    ) = formic::xdotc(n, &v.at(iter).at(0), 1, &Hv.at(i   ).at(0), 1);
          newS.at( i    + p*iter ) = formic::xdotc(n, &v.at(i   ).at(0), 1, &Sv.at(iter).at(0), 1);
          newS.at( iter + p*i    ) = formic::xdotc(n, &v.at(iter).at(0), 1, &Sv.at(i   ).at(0), 1);
          //newH[ p*iter + i ] = formic::conj(newH[ p*i + iter ]);
          //newS[ p*iter + i ] = formic::conj(newS[ p*i + iter ]);
        }

        // store the new subspace matrices
        subH = newH;
        subS = newS;

      }

      // diagonalize the subspace S matrix
      if (formic::mpi::rank() == 0) {
        formic::xcopy(p*p, &subS.at(0), 1, &s_evecs.at(0), 1);
        formic::xsyev('V', 'U', p, &s_evecs.at(0), p, &s_evals.at(0),
                      &work.at(0), lwork, &rwork.at(0), info);
        if ( info != 0 )
          throw formic::Exception("formic::xsyev failed for subspace S matrix with info = %i in formic::davidson") % info;
      }

      // get the smallest eigenvalue of the subspace S matrix
      double min_S_eval = *std::min_element(&s_evals.at(0), &s_evals.at(0) + p);
      formic::mpi::bcast(min_S_eval);

      // if the smallest eval is too small, stop iterating
      if ( min_S_eval < 1.0e-6 ) {
        if (formic::mpi::rank() == 0) {
          formic::of << boost::format("davidson iteration %4i:  stopping due to small subspace S eigenvalue of %.2e")
                        % iter
                        % min_S_eval
                     << std::endl;
        }
        break;
      }

      // initialize a flag to tell whether the lowest eigenvalue had any imaginary component
      bool need_to_symmetrize = symmetrize;

      // solve subspace eigenproblem on root process
      if (formic::mpi::rank() == 0) {

        // if requested, compute the (H,S) eigenvector/eigenvalue with a potentially non-hermitian H
        if ( ! symmetrize ) {

          // compute the product subS^(-1) subH
          formic::xgemm('T', 'N', p, p, p, formic::unity(T()), &s_evecs.at(0), p, &subH.at(0), p, formic::zero(T()), &matB.at(0), p);
          for (int j = 0; j < p; j++)
          for (int i = 0; i < p; i++)
            matB.at(i+p*j) /= s_evals.at(i);
          formic::xgemm('N', 'N', p, p, p, formic::unity(T()), &s_evecs.at(0), p, &matB.at(0), p, formic::zero(T()), &matA.at(0), p);

          // compute the eigenvalues and right eigenvectors of subS^(-1) subH
          formic::xgeev('N', 'V', p, &matA.at(0), p, &wc.at(0), &wr.at(0), &wi.at(0), &matB.at(0), p, &matB.at(0), p,
                        &work.at(0), lwork, &rwork.at(0), info);
          if ( info != 0 )
            throw formic::Exception("formic::xgeev failed for  subS^(-1) subH  with info = %i in formic::davidson") % info;

          // find the lowest eigenpair and make sure its eigenvalue is real
          double min_wr = wr.at(0);
          double min_wi = wi.at(0);
          int min_i = 0;
          for (int i = 0; i < p; i++) {
            if ( wr.at(i) < min_wr ) {
              min_wr = wr.at(i);
              min_wi = wi.at(i);
              min_i = i;
            }
          }

          // if the lowest eigenvalue is real, use it and its corresponding eigenvector
          if ( std::abs(min_wi) < 1.0e-10 ) {
            need_to_symmetrize = false;
            E = min_wr;
            for (int i = 0; i < p; i++)
             sub_evec.at(i) = matB.at( i + p * min_i );
          } else {
            need_to_symmetrize = true;
          }


        }

        // if requested, compute the (H,S) eigenvector/eigenvalue with a symmetrized H
        if ( need_to_symmetrize ) {

          // build the symmetrized subspace H and S matrices
          for (int j = 0; j < p; j++) {
            matA.at(j+p*j) = subH.at(j+p*j);
            matB.at(j+p*j) = subS.at(j+p*j);
            for (int i = 0; i < j; i++) {
              const T hval = 0.5 * ( formic::conj(subH.at(j+p*i)) + subH.at(i+p*j) );
              matA.at(i+p*j) = hval;
              matA.at(j+p*i) = formic::conj(hval);
              const T sval = 0.5 * ( formic::conj(subS.at(j+p*i)) + subS.at(i+p*j) );
              matB.at(i+p*j) = sval;
              matB.at(j+p*i) = formic::conj(sval);
            }
          }

          // compute the subspace eigenvectors and eigenvalues
          formic::xsygv(1, 'V', 'U', p, &matA.at(0), p, &matB.at(0), p, &wr.at(0), &work.at(0), lwork, &rwork.at(0), info);
          if ( info != 0 )
            throw formic::Exception("formic::xsygv failed for symmetrized subspace eigenproblem with info = %i in formic::davidson") % info;

          // extract the lowest eigenvalue and its corresponding eigenvector
          E = wr.at(0);
          for (int i = 0; i < p; i++)
            sub_evec.at(i) = matA.at(i+p*0);

        }

        // normalize the eigenvector
        T evec_norm = formic::zero(T());
        for (int i = 0; i < p; i++)
        for (int j = 0; j < p; j++)
          evec_norm += formic::conj(sub_evec.at(i)) * sub_evec.at(j) * subS.at(i+p*j);
        formic::xscal(p, formic::unity(T()) / std::abs(std::sqrt(evec_norm)), &sub_evec.at(0), 1);

      }

      // broadcast the eigenvalue to all processes
      formic::mpi::bcast(E);

      // construct the new residual
      std::fill(w.begin(), w.end(), formic::zero(T()));
      for (int i = 0; i < p; i++) {
        formic::xaxpy(n, sub_evec.at(i), &Hv.at(i).at(0), 1, &w.at(0), 1);
        formic::xaxpy(n, -E * sub_evec.at(i), &Sv.at(i).at(0), 1, &w.at(0), 1);
      }

      // compute the residual norm
      S->operate(n, &w.at(0), x, formic::mpi::rank(), 0);
      rnorm = std::abs(std::sqrt(formic::xdotc(n, &w.at(0), 1, x, 1)));

      // compute the new subspace basis vector
      v.at(p).assign(n, formic::zero(T()));
      if ( rnorm > 1.0e-1 ) // arnoldi when residual is large
        formic::xcopy(n, &w.at(0), 1, &v.at(p).at(0), 1);
      else // davidson when residual is small
        M->operate(n, &w.at(0), &v.at(p).at(0), formic::mpi::rank(), 0);

      // construct the new eigenvector
      std::fill(x, x+n, formic::zero(T()));
      for (int i = 0; i < p; i++)
        formic::xaxpy(n, sub_evec.at(i), &v.at(i).at(0), 1, x, 1);

      // save the new eigenvector
      formic::xcopy(n, x, 1, &best_x.at(0), 1);
      
      //// get eigenvector normalization
      //S->operate(n, x, &w.at(0), formic::mpi::rank(), 0);
      //const double enorm = std::abs(std::sqrt(formic::xdotc(n, x, 1, &w.at(0), 1)));

      // check if we are finished  (only the result on the root process is meaningful)
      converged = ( rnorm < thresh );
      //converged = ( std::abs(E - last_eval) < thresh );
      formic::mpi::bcast(converged);

      if (formic::mpi::rank() == 0)
        formic::of << boost::format("davidson iteration %4i:    eval = %20.12f          resid = %.2e      min_S_eval = %.2e%s%s")
                      % iter
                      % E
                      % rnorm
                      % min_S_eval
                      % ( need_to_symmetrize ? "    symm" : "" )
                      % ( E > last_eval      ? "    up" : "" )
                   << std::endl;

      // save this iteration's eigenvalue
      last_eval = E;

      //// print current solution vector
      //formic::of << std::endl;
      //formic::of << boost::format("solution vector %2i:") % iter << std::endl;
      //for (int i = 0; i < n; i++)
      //  formic::of << formic::format_number("%20.12f", x[i]) << std::endl;
      //formic::of << std::endl;

    }

    // check if the method converged
    if (converged && formic::mpi::rank() == 0)
      formic::of << boost::format("davidson solver converged in %10i iterations") % iter << std::endl << std::endl;
    else if (formic::mpi::rank() == 0)
      formic::of << boost::format("davidson solver did not converge after %10i iterations") % iter << std::endl << std::endl;
      //throw formic::Exception("formic::davidson did not converge.  residual = %.2e") % std::abs(rnorm);

    // place the best eigenvector in the output
    formic::xcopy(n, &best_x.at(0), 1, x, 1);

  }

} // end namespace formic

#endif
